import pandas as pd
import matplotlib.pyplot as plt
from brokenaxes import brokenaxes

plt.rcParams['font.sans-serif'] = ['Times New Roman']

styles = ['-', '-.', ':', '--', 'solid',
          'dashed', 'dotted', 'dashdot', 'dashed']
markers = [' ', '>', 'x', '*', '8', '+', 'p', 'D']
colors = ["red", "orange", "blue", "c", "cyan",
          "brown", "mediumvioletred", "dodgerblue", "green"]


def load_a_day_res(filename):
    """
    读取一个结果文件
    """
    with open(filename, 'r') as f:
        row_list = f.read().splitlines()
    a_day_res = row_list[0].split('|')
    a_day_res = [int(float(i)) for i in a_day_res]
    return a_day_res


def draw_plot(data_dict):
    """
    Input: 接受任意数量的数据，key充当画图的图例，value是画图用的原始数据
    """
    plt.figure(figsize=(8, 4))
    # bax = brokenaxes(ylims=((159000, 160000), (190000, 193000)), hspace=.05, despine=False)
    # 适配曲线数量
    count = 0
    for k, data in data_dict.items():
        x = []
        size = len(data)
        for i in range(size):
            x.append(i)
        # plt.bar(x, data, #label=k,
        #          linestyle=styles[count], color=colors[count], linewidth=2.5)

        # sns.lineplot(x, data, label=k, linestyle=styles[count], color=colors[count], linewidth=2.5)
        plt.plot(x, data, label=k, linestyle=styles[count], color=colors[count], linewidth=2.5, marker=markers[count], markersize=8)
        count += 1

    # plt.ylim(1000, 5000)
    # plt.xlim(0, 23)
    plt.yticks(fontsize=27)
    plt.xticks(fontsize=27)
    plt.xlabel("time(h)", fontsize=27)
    plt.ylabel("Energy consumption", fontsize=27)
    plt.grid()
    plt.legend(fontsize=27, loc='lower center', ncol=3, frameon=False, bbox_to_anchor=(0, 1.02, 1, 0.2),
        borderpad=0.1,handletextpad=0.1, handlelength=1, columnspacing=0.2, framealpha=0.5)
    return plt


if __name__ == "__main__":

    # energy_consumption = load_a_day_res('./energy_consumption/energy_1_0.2.txt')
    # baseline = [192000 for _ in range(len(energy_consumption))]
    origin = load_a_day_res('./energy_consumption/origin_1_0.2_2.txt')[0:168]
    method1 = load_a_day_res('./energy_consumption/method1_1_0.2_2.txt')[0:168]
    method2 = load_a_day_res('./energy_consumption/method2_1_0.2_2.txt')[0:168]
    
    data_dict = {
        'No-SS': origin,
        'SSwithoutSC': method1,
        'SSwithSC': method2,
    }

    pt = draw_plot(data_dict)
    pt.tight_layout()
    file_name = __file__.split('/')[-1].split('.')[0]
    pt.savefig(f"{file_name}.pdf", bbox_inches='tight')
    pt.show()
